/*
 * memory/main.c
 *
 * Copyright (C) 2018 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <memory.h>
#include <boot/multiboot.h>

#define MULTIBOOT_LOAD_ADDRESS 0x00100000
#define VIRTUAL_LOAD_ADDRESS   0x80000000

extern void memory_init_mapping_hack(void);
extern void memory_init_physical(memory_map_entry_t *mmap, size_t num_entries);
extern void (*_end)();

uintptr_t memory_metadata_base = MEMORY_METADATA_TOP;

static int compare_map_entry(const void *p1, const void *p2)
{
    memory_map_entry_t *a = (memory_map_entry_t*)p1;
    memory_map_entry_t *b = (memory_map_entry_t*)p2;

    if (a->address < b->address) return -1;
    else if (a->address > b->address) return 1;
    else return 0;
}

static multiboot_tag_t *early_map_mboot(uintptr_t mboot_tags, size_t mboot_size)
{
    multiboot_tag_t *mboot = (multiboot_tag_t*)(PAGE_ALIGN_UP((uintptr_t)&_end) + PAGE_OFFSET(mboot_tags));
    page_num_t num_pages = PAGE_NUMBER(mboot_tags + mboot_size - 1) - PAGE_NUMBER(mboot_tags) + 1;
    page_t pages[num_pages];

    area_t area = {
        .pages = pages,
        .count = num_pages,
    };

    for (page_num_t i = 0; i < num_pages; i++)
    {
        pages[i].status = PAGE_STATUS_ALLOCATED;
        pages[i].number = PAGE_NUMBER(mboot_tags) + i;
    }

    memory_map_area(memory_default_table, &area, mboot, MEMORY_FLAG_ACCESSIBLE);
    return mboot;
}

static size_t plot_memory_map(uintptr_t mboot_tags, size_t mboot_size, multiboot_tag_t *mboot, memory_map_entry_t *map, size_t max_entries)
{
    size_t num_entries = 0;

    if (num_entries >= max_entries) return 0;
    map[num_entries++] = (memory_map_entry_t) {
        .address = MULTIBOOT_LOAD_ADDRESS,
        .length = PAGE_ALIGN_UP((uintptr_t)&_end - VIRTUAL_LOAD_ADDRESS),
        .status = PAGE_STATUS_ALLOCATED,
    };

    if (num_entries >= max_entries) return 0;
    map[num_entries++] = (memory_map_entry_t) {
        .address = PAGE_ALIGN(mboot_tags),
        .length = PAGE_ALIGN_UP(mboot_tags + mboot_size) - PAGE_ALIGN(mboot_tags),
        .status = PAGE_STATUS_ALLOCATED,
    };

    if (num_entries >= max_entries) return 0;
    map[num_entries++] = (memory_map_entry_t) { .address = 0, .length = 0x1000, .status = PAGE_STATUS_RESERVED };

    for (multiboot_tag_t *tag = mboot; tag->type != MULTIBOOT_INFO_END; tag = (multiboot_tag_t*)(((uintptr_t)tag + tag->size + 7) & ~7))
    {
        switch (tag->type)
        {
        case MULTIBOOT_INFO_MEMORY_MAP:
            {
                multiboot_tag_mmap_t *mmap = (multiboot_tag_mmap_t*)tag;
                for (multiboot_mmap_entry_t *entry = (multiboot_mmap_entry_t*)(mmap + 1);
                     (uintptr_t)entry < ((uintptr_t)mmap + mmap->size);
                     entry = (multiboot_mmap_entry_t*)((uintptr_t)entry + mmap->entry_size))
                {
                    if (entry->type == 5) continue;
                    if (num_entries >= max_entries) return 0;
                    map[num_entries++] = (memory_map_entry_t) {
                        .address = PAGE_ALIGN(entry->base),
                        .length = PAGE_ALIGN_UP(entry->base + entry->length) - PAGE_ALIGN(entry->base),
                        .status = (entry->type == 1) ? PAGE_STATUS_FREE : PAGE_STATUS_RESERVED,
                    };
                }

                break;
            }

        case MULTIBOOT_INFO_MODULES:
            {
                multiboot_tag_module_t *module = (multiboot_tag_module_t*)tag;
                if (num_entries >= max_entries) return 0;
                map[num_entries++] = (memory_map_entry_t) {
                    .address = PAGE_ALIGN(module->mod_start),
                    .length = PAGE_ALIGN_UP(module->mod_end) - PAGE_ALIGN(module->mod_start),
                    .status = PAGE_STATUS_ALLOCATED,
                };

                break;
            }
        }
    }

    qsort(map, num_entries, sizeof(memory_map_entry_t), compare_map_entry);

    for (size_t i = 1; i < num_entries; i++)
    {
        uintptr_t current_end = map[i].address + map[i].length;
        uintptr_t previous_end = map[i - 1].address + map[i - 1].length;

        if (map[i].status > map[i - 1].status)
        {
            if (previous_end > current_end)
            {
                if (num_entries >= max_entries) return 0;
                memmove(&map[i + 2], &map[i + 1], (num_entries - i - 1) * sizeof(memory_map_entry_t));
                map[i + 1].address = current_end;
                map[i + 1].length = previous_end - current_end;
                map[i + 1].status = map[i - 1].status;
                num_entries++;
            }

            if (previous_end > map[i].address)
            {
                if (map[i - 1].length > previous_end - map[i].address)
                {
                    map[i - 1].length -= previous_end - map[i].address;
                }
                else
                {
                    memmove(&map[i - 1], &map[i], (num_entries - i) * sizeof(memory_map_entry_t));
                    i -= 2;
                    num_entries--;
                }
            }
        }
        else if (map[i].status < map[i - 1].status)
        {
            if (previous_end >= current_end)
            {
                memmove(&map[i], &map[i + 1], (num_entries - i - 1) * sizeof(memory_map_entry_t));
                i--;
                num_entries--;
            }
            else if (previous_end > map[i].address)
            {
                map[i].length -= previous_end - map[i].address;
                map[i].address = previous_end;
            }
        }
        else if (map[i].address <= previous_end)
        {
            map[i - 1].length = (current_end > previous_end ? current_end : previous_end) - map[i - 1].address;
            memmove(&map[i], &map[i + 1], (num_entries - i - 1) * sizeof(memory_map_entry_t));
            num_entries--;
        }
    }

    return num_entries;
}

void *memory_request_metadata_space(size_t count, size_t size)
{
    memory_metadata_base -= count * size;
    if (memory_metadata_base % size) memory_metadata_base = size * (memory_metadata_base / size);
    return (void*)memory_metadata_base;
}

void memory_init(uintptr_t mboot_tags, size_t mboot_size)
{
    cpu_max_physical_bits = 32; // No PAE for now
    size_t num_entries, max_entries = 1024;
    multiboot_tag_t *mboot = early_map_mboot(mboot_tags, mboot_size);

    for (;;)
    {
        memory_map_entry_t map[max_entries];

        if ((num_entries = plot_memory_map(mboot_tags, mboot_size, mboot, map, max_entries)))
        {
            memory_init_physical(map, num_entries);
            break;
        }

        max_entries += max_entries;
    }

    area_t kernel_area = {
        .pages = memory_find_page_by_address(MULTIBOOT_LOAD_ADDRESS),
        .count = PAGE_NUMBER(PAGE_ALIGN_UP((uintptr_t)&_end - VIRTUAL_LOAD_ADDRESS)),
    };

    ASSERT(kernel_area.pages != NULL);
    memory_init_virtual(&kernel_area);
}
